#include <linux/net.h>
#include <linux/in.h>
#include <linux/fs.h>
#include "kblock.h"
extern struct kblock_device *kdevice;

static int kblock_recv(struct socket *sock,char* buf,int size);
static int kblock_recv_all(struct socket *sock,char *buf,int size);

static int kblock_recv_all(struct socket *sock,char *buf,int size)
{
    int left = size;
    char *ptr = buf;
    int recvlen;

    while(left) {
        if (atomic_read(&kdevice->status) == KBLOCK_EXIT) {
            return -1;
        }
        recvlen = kblock_recv(sock,ptr,left);
        if (recvlen == -EINTR) {
            printk(KERN_INFO "%s:kblock_recv interrupted\n",__FUNCTION__);
            continue;
        }
        else if (recvlen == -EAGAIN) {
            //printk(KERN_INFO "%s:kblock_recv EAGAIN\n",__FUNCTION__);
            continue;
        }else if (recvlen < 0 ) {
            printk(KERN_INFO "%s:kblock_recv failed\n",__FUNCTION__);
            return -1;
        }
        else if(recvlen == 0) {
            printk(KERN_INFO "%s:sock is closed by peer",__FUNCTION__);
            return -1;
        }
        ptr += recvlen;
        left -= recvlen;
    }
    return 0;
}

static int kblock_recv(struct socket *sock,char* buf,int size)
{
    struct msghdr msg = {0};
    struct kvec vec = {0};
    int error;

    vec.iov_base = buf;
    vec.iov_len = size;

    msg.msg_iovlen = 1;
    msg.msg_iov = (struct iovec*)&vec;
    msg.msg_flags = MSG_WAITALL | MSG_NOSIGNAL;

    error = kernel_recvmsg(sock,&msg,&vec,1,size,0);
    return error;
}

int kblock_recv_bio(struct kblock_device *device)
{
    struct data_head head = {0};
    struct bio *bio = NULL;
    unsigned int data_len = 0;
    sector_t sector = 0;
    int type = 0;
    unsigned int npage = 0;
    int i = 0;
    void* addr = NULL;
    struct page *page = NULL;

    printk(KERN_INFO "kblock_recv_bio:wait for recv head\n");
    if (kblock_recv_all(device->data_sock,(char *)&head,sizeof(struct data_head)) < 0) {
        printk(KERN_INFO "kblock_recv_bio:recv head failed\n");
        return -1;
    }

    if (KBLOCK_MAGIC != be32_to_cpu(head.magic)) {
        printk(KERN_ERR "bad data head\n");
        return -1;
    }
    type = be32_to_cpu(head.type);
    data_len = be32_to_cpu(head.data_len);
    sector = be64_to_cpu(head.sector);
    
    npage = (data_len + PAGE_SIZE -1) / PAGE_SIZE;
    printk(KERN_INFO "bio page number %d\n",npage);
    bio = bio_alloc(GFP_NOIO,npage);
    if (bio == NULL) {
        printk(KERN_INFO "bio_alloc failed\n");
        return -1;
    }

    bio->bi_bdev = device->physics_device;
    bio->bi_sector = sector;
    bio->bi_private = device;
    bio->bi_rw = WRITE;
    bio->bi_end_io = kblock_bio_endio;

    for (i = 0;i < npage;i++) {
        page = alloc_page(GFP_KERNEL); 
        if (page == NULL) {
            printk(KERN_INFO "alloc page failed\n");
            return -1;
        }
        addr = kmap(page);
        BUG_ON(addr == NULL);

        if (kblock_recv_all(device->data_sock,addr,data_len) < 0) {
            printk(KERN_INFO "kblock_recv_bio:recv page failed\n");
            goto failed;
        }
        kunmap(page);
        bio_add_page(bio,page,data_len,0);
    }
    generic_make_request(bio);
    return 0;
failed:
    printk(KERN_INFO "kblock_recv_bio failed\n");
    /* need to free alloc memory */
    return -1;
}

